"""
# Turing Pattern Formation - Research Application

This module demonstrates reaction-diffusion systems that generate Turing patterns,
including spots, stripes, and complex spatial structures. Based on the original
work by Alan Turing (1952) on morphogenesis.

## Mathematical Background

The Turing system consists of two coupled reaction-diffusion equations:
```
∂u/∂t = Du ∇²u + f(u,v)
∂v/∂t = Dv ∇²v + g(u,v)
```

Where:
- u: Activator species concentration
- v: Inhibitor species concentration
- Du, Dv: Diffusion coefficients (typically Dv > Du)
- f(u,v), g(u,v): Reaction kinetics

## Implementation Features

- Multiple reaction kinetics (FitzHugh-Nagumo, Schnakenberg, Gray-Scott)
- Automatic stability analysis and pattern prediction
- Parameter space exploration tools
- Publication-quality visualization
- Performance optimization for large-scale simulations

## References

- Turing, A. M. (1952). The chemical basis of morphogenesis. Phil. Trans. R. Soc. B 237, 37-72.
- Murray, J. D. (2003). Mathematical Biology II: Spatial Models and Biomedical Applications.
- Cross, M. C. & Hohenberg, P. C. (1993). Pattern formation outside of equilibrium. Rev. Mod. Phys. 65, 851.
"""

using LinearAlgebra, FFTW, Plots, Statistics
using Printf, BenchmarkTools

# Export main functions
export TuringSystem, run_turing_simulation, analyze_patterns, create_pattern_gallery

"""
    TuringSystem

Represents a two-species reaction-diffusion system for Turing pattern formation.

# Fields
- `Du::Float64`: Diffusion coefficient for activator u
- `Dv::Float64`: Diffusion coefficient for inhibitor v  
- `kinetics::Symbol`: Type of reaction kinetics (:fitzhugh_nagumo, :schnakenberg, :gray_scott)
- `parameters::NamedTuple`: Kinetic parameters specific to the chosen model
- `domain_size::Tuple`: Spatial domain dimensions (Lx, Ly)
- `grid_size::Tuple`: Number of grid points (nx, ny)
"""
struct TuringSystem
    Du::Float64
    Dv::Float64
    kinetics::Symbol
    parameters::NamedTuple
    domain_size::Tuple{Float64, Float64}
    grid_size::Tuple{Int, Int}
end

"""
    TuringSystem(kinetics::Symbol; Du=1.0, Dv=10.0, domain_size=(20.0, 20.0), grid_size=(128, 128), kwargs...)

Create a Turing system with specified kinetics and parameters.

# Examples
```julia
# FitzHugh-Nagumo system
sys = TuringSystem(:fitzhugh_nagumo, a=0.1, b=0.9, τ=1.0)

# Schnakenberg system  
sys = TuringSystem(:schnakenberg, a=0.1, b=0.9)

# Gray-Scott system
sys = TuringSystem(:gray_scott, F=0.04, k=0.06)
```
"""
function TuringSystem(kinetics::Symbol; 
                     Du=1.0, Dv=10.0, 
                     domain_size=(20.0, 20.0), 
                     grid_size=(128, 128),
                     kwargs...)
    # Set default parameters for each kinetic model
    if kinetics == :fitzhugh_nagumo
        default_params = (a=0.1, b=0.9, τ=1.0)
    elseif kinetics == :schnakenberg
        default_params = (a=0.1, b=0.9)
    elseif kinetics == :gray_scott
        default_params = (F=0.04, k=0.06)
    else
        throw(ArgumentError("Unknown kinetics: $kinetics"))
    end
    
    # Merge with user-provided parameters
    parameters = merge(default_params, NamedTuple(kwargs))
    
    return TuringSystem(Du, Dv, kinetics, parameters, domain_size, grid_size)
end

"""
    reaction_kinetics(u, v, system::TuringSystem)

Compute reaction terms f(u,v) and g(u,v) for the specified kinetic model.
"""
function reaction_kinetics(u, v, system::TuringSystem)
    if system.kinetics == :fitzhugh_nagumo
        p = system.parameters
        fu = u - u^3 - v + p.a
        gv = (u - v) / p.τ
        return fu, gv
        
    elseif system.kinetics == :schnakenberg
        p = system.parameters
        fu = p.a - u + u^2 * v
        gv = p.b - u^2 * v
        return fu, gv
        
    elseif system.kinetics == :gray_scott
        p = system.parameters
        fu = -u * v^2 + p.F * (1 - u)
        gv = u * v^2 - (p.F + p.k) * v
        return fu, gv
    end
end

"""
    initialize_system(system::TuringSystem, initial_condition::Symbol=:random_perturbation)

Initialize the concentration fields u and v with appropriate initial conditions.

# Arguments
- `system::TuringSystem`: The Turing system configuration
- `initial_condition::Symbol`: Type of initial condition
  - `:random_perturbation`: Small random perturbations around steady state
  - `:localized_spots`: Localized high-concentration regions
  - `:stripes`: Initial stripe-like pattern
"""
function initialize_system(system::TuringSystem, initial_condition::Symbol=:random_perturbation)
    nx, ny = system.grid_size
    Lx, Ly = system.domain_size
    
    # Create coordinate grids
    x = range(0, Lx, length=nx)
    y = range(0, Ly, length=ny)
    X = [xi for xi in x, yi in y]
    Y = [yi for xi in x, yi in y]
    
    # Find steady state for the given kinetics
    u_steady, v_steady = find_steady_state(system)
    
    if initial_condition == :random_perturbation
        # Small random perturbations around steady state
        u = u_steady .+ 0.01 * randn(nx, ny)
        v = v_steady .+ 0.01 * randn(nx, ny)
        
    elseif initial_condition == :localized_spots
        # Localized high-concentration spots
        u = u_steady * ones(nx, ny)
        v = v_steady * ones(nx, ny)
        
        # Add localized perturbations
        for _ in 1:3
            cx, cy = Lx * rand(), Ly * rand()
            σ = min(Lx, Ly) / 20
            u .+= 0.1 * exp.(-((X .- cx).^2 + (Y .- cy).^2) / (2σ^2))
        end
        
    elseif initial_condition == :stripes
        # Initial stripe pattern
        k_stripe = 2π / (min(Lx, Ly) / 4)  # Wavelength
        u = u_steady .+ 0.1 * cos.(k_stripe * X)
        v = v_steady .+ 0.1 * cos.(k_stripe * X + π/2)
        
    else
        throw(ArgumentError("Unknown initial condition: $initial_condition"))
    end
    
    return u, v, X, Y
end

"""
    find_steady_state(system::TuringSystem)

Find the homogeneous steady state (u₀, v₀) where f(u₀,v₀) = g(u₀,v₀) = 0.
"""
function find_steady_state(system::TuringSystem)
    if system.kinetics == :fitzhugh_nagumo
        # For FitzHugh-Nagumo: u₀ - u₀³ - v₀ + a = 0, (u₀ - v₀)/τ = 0
        # This gives v₀ = u₀ and u₀ - u₀³ + a = 0
        p = system.parameters
        u0 = solve_cubic(1.0, 0.0, 1.0, p.a)  # Approximate solution
        v0 = u0
        
    elseif system.kinetics == :schnakenberg
        # For Schnakenberg: a - u₀ + u₀²v₀ = 0, b - u₀²v₀ = 0
        # From second equation: u₀²v₀ = b
        # Substituting: a - u₀ + b = 0, so u₀ = a + b
        p = system.parameters
        u0 = p.a + p.b
        v0 = p.b / u0^2
        
    elseif system.kinetics == :gray_scott
        # For Gray-Scott: -u₀v₀² + F(1-u₀) = 0, u₀v₀² - (F+k)v₀ = 0
        # Steady state is typically (1, 0) for the parameter regime of interest
        u0 = 1.0
        v0 = 0.0
    end
    
    return u0, v0
end

"""
    solve_cubic(a, b, c, d)

Solve cubic equation ax³ + bx² + cx + d = 0, returning the real root nearest to 0.
"""
function solve_cubic(a, b, c, d)
    # Simplified solver for the specific case in FitzHugh-Nagumo
    # x - x³ + a = 0, so x³ - x - a = 0
    # Use Newton's method starting from x = a
    x = a
    for _ in 1:10
        f = x^3 - x - a
        df = 3x^2 - 1
        x_new = x - f/df
        if abs(x_new - x) < 1e-12
            break
        end
        x = x_new
    end
    return x
end

"""
    evolve_system!(u, v, u_new, v_new, system::TuringSystem, dt, dx, dy)

Perform one time step evolution using finite differences.
"""
function evolve_system!(u, v, u_new, v_new, system::TuringSystem, dt, dx, dy)
    nx, ny = size(u)
    Du, Dv = system.Du, system.Dv
    
    # Interior points using finite differences
    for j in 2:ny-1, i in 2:nx-1
        # Laplacian using 5-point stencil
        ∇²u = (u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1] - 4u[i,j]) / dx^2
        ∇²v = (v[i-1,j] + v[i+1,j] + v[i,j-1] + v[i,j+1] - 4v[i,j]) / dx^2
        
        # Reaction terms
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        
        # Time evolution
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
    end
    
    # Periodic boundary conditions
    # Left-right boundaries
    for j in 2:ny-1
        i = 1
        ∇²u = (u[nx-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1] - 4u[i,j]) / dx^2
        ∇²v = (v[nx-1,j] + v[i+1,j] + v[i,j-1] + v[i,j+1] - 4v[i,j]) / dx^2
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
        
        i = nx
        ∇²u = (u[i-1,j] + u[2,j] + u[i,j-1] + u[i,j+1] - 4u[i,j]) / dx^2
        ∇²v = (v[i-1,j] + v[2,j] + v[i,j-1] + v[i,j+1] - 4v[i,j]) / dx^2
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
    end
    
    # Top-bottom boundaries
    for i in 2:nx-1
        j = 1
        ∇²u = (u[i-1,j] + u[i+1,j] + u[i,ny-1] + u[i,j+1] - 4u[i,j]) / dx^2
        ∇²v = (v[i-1,j] + v[i+1,j] + v[i,ny-1] + v[i,j+1] - 4v[i,j]) / dx^2
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
        
        j = ny
        ∇²u = (u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,2] - 4u[i,j]) / dx^2
        ∇²v = (v[i-1,j] + v[i+1,j] + v[i,j-1] + v[i,2] - 4v[i,j]) / dx^2
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
    end
    
    # Corner points
    corners = [(1,1), (1,ny), (nx,1), (nx,ny)]
    for (i,j) in corners
        # Simplified treatment for corners
        i_prev = i == 1 ? nx-1 : i-1
        i_next = i == nx ? 2 : i+1
        j_prev = j == 1 ? ny-1 : j-1
        j_next = j == ny ? 2 : j+1
        
        ∇²u = (u[i_prev,j] + u[i_next,j] + u[i,j_prev] + u[i,j_next] - 4u[i,j]) / dx^2
        ∇²v = (v[i_prev,j] + v[i_next,j] + v[i,j_prev] + v[i,j_next] - 4v[i,j]) / dx^2
        fu, gv = reaction_kinetics(u[i,j], v[i,j], system)
        u_new[i,j] = u[i,j] + dt * (Du * ∇²u + fu)
        v_new[i,j] = v[i,j] + dt * (Dv * ∇²v + gv)
    end
end

"""
    run_turing_simulation(system::TuringSystem; 
                         t_final=100.0, 
                         dt=nothing, 
                         initial_condition=:random_perturbation,
                         save_interval=1.0,
                         verbose=true)

Run a complete Turing pattern formation simulation.

# Returns
- `NamedTuple` with simulation results including time series, final patterns, and diagnostics
"""
function run_turing_simulation(system::TuringSystem; 
                               t_final=100.0, 
                               dt=nothing, 
                               initial_condition=:random_perturbation,
                               save_interval=1.0,
                               verbose=true)
    
    if verbose
        println("🧬 Starting Turing Pattern Simulation")
        println("====================================")
        println("Kinetics: $(system.kinetics)")
        println("Domain: $(system.domain_size)")
        println("Grid: $(system.grid_size)")
        println("Diffusion ratio Dv/Du: $(system.Dv/system.Du)")
        println()
    end
    
    # Initialize system
    u, v, X, Y = initialize_system(system, initial_condition)
    nx, ny = system.grid_size
    Lx, Ly = system.domain_size
    dx = Lx / nx
    dy = Ly / ny
    
    # Determine stable time step
    if dt === nothing
        dt = min(dx^2, dy^2) / (4 * max(system.Du, system.Dv)) * 0.8
    end
    
    n_steps = Int(t_final / dt)
    n_save = max(1, Int(save_interval / dt))
    
    if verbose
        println("Time integration:")
        println("  dt = $(round(dt, digits=6))")
        println("  Steps: $n_steps")
        println("  Save every: $n_save steps")
        println()
    end
    
    # Storage for time series
    times = Float64[]
    u_snapshots = Array{Float64,3}(undef, nx, ny, 0)
    v_snapshots = Array{Float64,3}(undef, nx, ny, 0)
    
    # Pre-allocate work arrays
    u_new = similar(u)
    v_new = similar(v)
    
    # Time evolution
    t = 0.0
    step = 0
    
    # Save initial condition
    push!(times, t)
    u_snapshots = cat(u_snapshots, reshape(u, nx, ny, 1), dims=3)
    v_snapshots = cat(v_snapshots, reshape(v, nx, ny, 1), dims=3)
    
    if verbose
        print("Progress: ")
    end
    
    while t < t_final
        step += 1
        
        # Time step
        evolve_system!(u, v, u_new, v_new, system, dt, dx, dy)
        u, u_new = u_new, u
        v, v_new = v_new, v
        t += dt
        
        # Save snapshots
        if step % n_save == 0
            push!(times, t)
            u_snapshots = cat(u_snapshots, reshape(u, nx, ny, 1), dims=3)
            v_snapshots = cat(v_snapshots, reshape(v, nx, ny, 1), dims=3)
            
            if verbose
                progress = t / t_final * 100
                print("$(Int(round(progress)))% ")
            end
        end
    end
    
    if verbose
        println("\n✅ Simulation completed!")
        println()
    end
    
    # Return comprehensive results
    return (
        times = times,
        u_final = u,
        v_final = v,
        u_snapshots = u_snapshots,
        v_snapshots = v_snapshots,
        X = X,
        Y = Y,
        system = system,
        dt = dt,
        n_steps = n_steps
    )
end

"""
    analyze_patterns(results; verbose=true)

Analyze the patterns formed in a Turing simulation.

# Returns
- Pattern analysis including wavelengths, correlation functions, and stability metrics
"""
function analyze_patterns(results; verbose=true)
    u_final = results.u_final
    v_final = results.v_final
    X, Y = results.X, results.Y
    system = results.system
    
    if verbose
        println("📊 Pattern Analysis")
        println("===================")
    end
    
    # Basic statistics
    u_mean = mean(u_final)
    u_std = std(u_final)
    v_mean = mean(v_final)
    v_std = std(v_final)
    
    # Pattern intensity (deviation from mean)
    pattern_intensity_u = std(u_final) / mean(u_final)
    pattern_intensity_v = std(v_final) / mean(v_final)
    
    if verbose
        println("Concentration statistics:")
        println("  u: $(round(u_mean, digits=4)) ± $(round(u_std, digits=4))")
        println("  v: $(round(v_mean, digits=4)) ± $(round(v_std, digits=4))")
        println("Pattern intensity:")
        println("  u: $(round(pattern_intensity_u, digits=4))")
        println("  v: $(round(pattern_intensity_v, digits=4))")
    end
    
    # Spatial correlation analysis via FFT
    u_centered = u_final .- u_mean
    u_fft = fft(u_centered)
    power_spectrum = abs.(u_fft).^2
    
    # Find dominant wavelength
    nx, ny = size(u_final)
    kx = fftfreq(nx, nx/system.domain_size[1])
    ky = fftfreq(ny, ny/system.domain_size[2])
    k_magnitude = sqrt.(kx'.^2 .+ ky.^2)
    
    # Radial average of power spectrum
    k_max = minimum(size(power_spectrum)) ÷ 2
    k_bins = range(0, k_max * minimum(abs.([kx[2]-kx[1], ky[2]-ky[1]])), length=k_max)
    power_radial = zeros(length(k_bins)-1)
    
    for i in 1:length(k_bins)-1
        mask = (k_magnitude .>= k_bins[i]) .& (k_magnitude .< k_bins[i+1])
        if any(mask)
            power_radial[i] = mean(power_spectrum[mask])
        end
    end
    
    # Find peak wavelength
    peak_idx = argmax(power_radial[2:end]) + 1  # Skip DC component
    dominant_k = (k_bins[peak_idx] + k_bins[peak_idx+1]) / 2
    dominant_wavelength = 2π / dominant_k
    
    if verbose
        println("Spatial analysis:")
        println("  Dominant wavelength: $(round(dominant_wavelength, digits=2))")
        println("  Dominant wavenumber: $(round(dominant_k, digits=4))")
    end
    
    # Pattern classification based on structure
    # Analyze isotropy by comparing power along different directions
    power_x = mean(abs.(fft(mean(u_centered, dims=2)[:, 1])).^2)
    power_y = mean(abs.(fft(mean(u_centered, dims=1)[1, :])).^2)
    anisotropy = abs(power_x - power_y) / (power_x + power_y)
    
    pattern_type = if anisotropy > 0.5
        "stripes"
    elseif pattern_intensity_u > 0.1
        "spots"
    else
        "weak_pattern"
    end
    
    if verbose
        println("  Pattern type: $pattern_type")
        println("  Anisotropy: $(round(anisotropy, digits=3))")
        println()
    end
    
    return (
        u_statistics = (mean=u_mean, std=u_std, intensity=pattern_intensity_u),
        v_statistics = (mean=v_mean, std=v_std, intensity=pattern_intensity_v),
        dominant_wavelength = dominant_wavelength,
        dominant_wavenumber = dominant_k,
        power_spectrum = power_spectrum,
        power_radial = power_radial,
        k_bins = k_bins,
        anisotropy = anisotropy,
        pattern_type = pattern_type
    )
end

"""
    create_pattern_gallery(results; analysis=nothing, save_path=nothing)

Create publication-quality visualizations of Turing patterns.
"""
function create_pattern_gallery(results; analysis=nothing, save_path=nothing)
    println("🎨 Creating Pattern Gallery")
    println("==========================")
    
    u_final = results.u_final
    v_final = results.v_final
    times = results.times
    
    # Create multi-panel figure
    p1 = heatmap(u_final', title="Activator (u)", aspect_ratio=1, 
                 color=:viridis, xlabel="x", ylabel="y")
    
    p2 = heatmap(v_final', title="Inhibitor (v)", aspect_ratio=1,
                 color=:plasma, xlabel="x", ylabel="y")
    
    # Time evolution of snapshots (if available)
    if length(times) > 3
        n_snapshots = min(4, length(times))
        snapshot_indices = Int.(round.(range(1, length(times), length=n_snapshots)))
        
        plots_evolution = []
        for (i, idx) in enumerate(snapshot_indices)
            t = times[idx]
            u_snap = results.u_snapshots[:, :, idx]
            p = heatmap(u_snap', title="t = $(round(t, digits=1))", 
                       aspect_ratio=1, color=:viridis, showaxis=false)
            push!(plots_evolution, p)
        end
        
        p_evolution = plot(plots_evolution..., layout=(1, n_snapshots), 
                          plot_title="Pattern Evolution")
    else
        p_evolution = plot(title="Evolution data not available")
    end
    
    # Combine all plots
    main_plot = plot(p1, p2, layout=(1, 2), size=(800, 300))
    full_plot = plot(main_plot, p_evolution, layout=(2, 1), size=(800, 600))
    
    if save_path !== nothing
        savefig(full_plot, save_path)
        println("Gallery saved to: $save_path")
    end
    
    println("Gallery created successfully!")
    return full_plot
end

"""
    parameter_space_study(kinetics::Symbol; parameter_ranges..., grid_points=5)

Explore parameter space for pattern formation boundaries.

# Example
```julia
results = parameter_space_study(:schnakenberg, 
                               a = 0.05:0.05:0.2, 
                               b = 0.8:0.1:1.2)
```
"""
function parameter_space_study(kinetics::Symbol; parameter_ranges..., grid_points=5)
    println("🔍 Parameter Space Exploration")
    println("==============================")
    
    param_dict = Dict(parameter_ranges)
    param_names = collect(keys(param_dict))
    param_values = collect(values(param_dict))
    
    println("Parameters to explore:")
    for (name, values) in param_dict
        println("  $name: $(first(values)) to $(last(values)) ($(length(values)) points)")
    end
    println()
    
    results_matrix = Dict()
    total_runs = prod(length.(param_values))
    run_count = 0
    
    for param_combination in Iterators.product(param_values...)
        run_count += 1
        
        # Create parameter NamedTuple
        params = NamedTuple{tuple(param_names...)}(param_combination)
        
        print("Run $run_count/$total_runs: ")
        for (name, val) in zip(param_names, param_combination)
            print("$name=$(round(val, digits=3)) ")
        end
        
        try
            # Create system with current parameters
            system = TuringSystem(kinetics; params...)
            
            # Run short simulation to check for pattern formation
            result = run_turing_simulation(system, t_final=50.0, verbose=false)
            analysis = analyze_patterns(result, verbose=false)
            
            # Store results
            results_matrix[param_combination] = (
                system = system,
                result = result,
                analysis = analysis,
                success = true
            )
            
            println("✅ $(analysis.pattern_type)")
            
        catch e
            println("❌ Failed: $e")
            results_matrix[param_combination] = (
                system = nothing,
                result = nothing, 
                analysis = nothing,
                success = false,
                error = e
            )
        end
    end
    
    println("\nParameter space exploration completed!")
    return results_matrix
end

# Example usage and demonstration
if abspath(PROGRAM_FILE) == @__FILE__
    println("🧬 Turing Pattern Formation - Research Example")
    println("=" * 50)
    
    # Example 1: FitzHugh-Nagumo system
    println("\n1. FitzHugh-Nagumo System")
    system_fn = TuringSystem(:fitzhugh_nagumo, Du=1.0, Dv=10.0, a=0.1, b=0.9)
    results_fn = run_turing_simulation(system_fn, t_final=50.0)
    analysis_fn = analyze_patterns(results_fn)
    
    # Example 2: Schnakenberg system  
    println("\n2. Schnakenberg System")
    system_sch = TuringSystem(:schnakenberg, Du=0.05, Dv=1.0, a=0.1, b=0.9)
    results_sch = run_turing_simulation(system_sch, t_final=100.0)
    analysis_sch = analyze_patterns(results_sch)
    
    # Example 3: Gray-Scott system
    println("\n3. Gray-Scott System")  
    system_gs = TuringSystem(:gray_scott, Du=0.16, Dv=0.08, F=0.04, k=0.06)
    results_gs = run_turing_simulation(system_gs, t_final=100.0, 
                                     initial_condition=:localized_spots)
    analysis_gs = analyze_patterns(results_gs)
    
    println("\n🎉 All examples completed successfully!")
    println("Ready for research applications and further exploration.")
end